# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question2/hydrogen.py
import math

import numpy
import scipy


class Hydrogen:
    def __init__(self,
                 assoc_laguerre_function=scipy.special.assoc_laguerre,
                 sph_harm_function=scipy.special.sph_harm):
        self.assoc_laguerre = assoc_laguerre_function
        self.sph_harm = sph_harm_function

        self.n = 1
        self.l_ = 0
        self.m = 0

    def set_state(self, n=None, l_=None, m=None):
        if n:
            self.n = n
        if m:
            self.m = m
        if l_:
            self.l_ = l_

    def radial(self, r):
        a0 = scipy.constants.value("Bohr radius")
        n = self.n
        l_ = self.l_

        t = r / a0

        result = (2 / n / a0) ** (3 / 2)
        result *= (math.factorial(n - l_ - 1) / 2 / n / math.factorial(n + l_)) ** (1 / 2)
        result *= (2 * t / n) ** l_
        result *= numpy.exp(-t / n)
        result *= self.assoc_laguerre(2 * t / n, n - l_ - 1, 2 * l_ + 1)

        return result

    # noinspection SpellCheckingInspection
    def poloidal(self, theta, phi):
        return self.sph_harm(self.m, self.l_, theta, phi)
